import pickle
import pandas as pd
from dayCalculation import *
from xgboost import XGBClassifier
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
from xgboost import plot_importance
from sklearn.metrics import accuracy_score  


def train_model(stocks_info):
    trade_days = list(stocks_info[list(stocks_info.keys())[0]].keys())
    Y = []
    X = []
    #选几个显著的因子作为特征
    index = [23, 21, 24, 33, 6, 32, 26, 19, 25, 18, 27, 20]
    for stock in stocks_info:
        for day in stocks_info[stock]:
            #使用2010-2014年的数据进行训练
            if day.split('-')[0] not in ['2010', '2011', '2012', '2013', '2014']:
                continue
            #用30天后的涨跌情况作为label进行二分类
            sell_day = next_day(day, 30)
            while sell_day not in trade_days and isLater(sell_day, trade_days[-1]):
                sell_day = next_day(sell_day, 1)
            if not isLater(sell_day, trade_days[-1]):
                break
            y = 1 if stocks_info[stock][sell_day]['close'] - stocks_info[stock][day]['close'] > 0 else 0
            x = [stocks_info[stock][day].values[i] for i in index]
            Y.append(y)
            X.append(x)

    x_train, x_test, y_train, y_test = train_test_split(X, Y, test_size=0.3)
    # XGBoost启动！
    model = XGBClassifier()
    model.fit(x_train, y_train)
    y_pred = model.predict(x_test)
    accuracy = accuracy_score(y_test, y_pred)
    #正确率83%左右，还可以
    print("machine learning accuarcy: %.2f%%" % (accuracy * 100.0))
    '''
    #看看哪些特征最显著
    fig,ax = plt.subplots(figsize=(10,15))
    plot_importance(model,height=0.5,max_num_features=64,ax=ax)
    plt.show()
    '''
    return model


if __name__ == '__main__':
    with open('data.pkl', 'rb') as f:
        data = pickle.load(f)
    stocks_info = {}
    for day in data:
        df = data[day]
        for index, row in df.iterrows():
            if index not in stocks_info:
                stocks_info[index] = {}
            stocks_info[index][day] = row
    train_model(stocks_info)
